﻿// this file is copyied from megui under GPLv2
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace AutoEncodeDaemon
{
    public class Size
    {
        public Size(int width, int height)
        {
            Width = width;
            Height = height;
        }
        public int Width;
        public int Height;
    }
    public enum Unit : ushort { B = 0, KB, MB, GB }
    public enum ContainerType { MP4, MKV, AVI, M2TS };
    public enum AudioType { NONE,AC3, MP4AAC, VBRMP3, M4A, RAWAAC, CBRMP3, MP3, DTS, VORBIS, DTSHD, THD, PCM, DTSMA };
    public class VideoCodec
    {
        private string id;
        public string ID
        {
            get { return id; }
        }
        public VideoCodec(string id)
        {
            this.id = id;
        }
        public static readonly VideoCodec ASP = new VideoCodec("ASP");
        public static readonly VideoCodec AVC = new VideoCodec("AVC");
        public static readonly VideoCodec HFYU = new VideoCodec("HFYU");
        public static readonly VideoCodec VC1 = new VideoCodec("VC1");
        public static readonly VideoCodec MPEG2 = new VideoCodec("MPEG2");
    }
    public struct FileSize
    {
        public static readonly FileSize Empty = new FileSize(0);
        public static readonly FileSize MinNonZero = new FileSize(1);


        public static FileSize Of(string filename)
        {
            FileInfo fi = new FileInfo(filename);
            return new FileSize(fi.Length);
        }

        public static FileSize? Of2(string filename)
        {
            if (string.IsNullOrEmpty(filename) || !File.Exists(filename)) return null;

            FileInfo fi = new FileInfo(filename);
            if (fi.Length <= 0) return null;

            return new FileSize(fi.Length);
        }
        public static readonly string[] UnitSuffixes = new string[] { "bytes", "KB", "MB", "GB" };
        private ulong sizeInBytes;

        #region helper functions
        /// <summary>
        /// Returns the stored size, divided by divisor and rounded appropriately
        /// </summary>
        /// <param name="divisor"></param>
        /// <returns></returns>
        private ulong divideAndRound(ulong divisor)
        {
            ulong divided = sizeInBytes / divisor;
            if (sizeInBytes % divisor > (divisor / 2))
                divided++;
            return divided;
        }

        /// <summary>
        /// Sets sizeInBytes to (value * divisor) rounded appropriately
        /// </summary>
        /// <param name="divisor"></param>
        /// <param name="value"></param>
        private void setExact(ulong divisor, decimal value)
        {
            decimal result = value * ((decimal)divisor);
            sizeInBytes = (ulong)Math.Round(result);
        }
        #endregion

        #region generic Unit accessors
        private ulong divisor(Unit u)
        {
            ulong divisor = 1;
            for (ushort i = 0; i < (ushort)u; i++)
                divisor *= 1024;
            return divisor;
        }

        /// <summary>
        /// Returns the stored size in the requested units (rounded)
        /// </summary>
        /// <param name="shift"></param>
        /// <returns></returns>
        public ulong InUnits(Unit u)
        {
            return divideAndRound(divisor(u));
        }

        /// <summary>
        /// Returns the stored size in the requested units (exact)
        /// </summary>
        /// <param name="u"></param>
        /// <returns></returns>
        public decimal InUnitsExact(Unit u)
        {
            return ((decimal)sizeInBytes) / ((decimal)divisor(u));
        }

        public void SetWithUnits(Unit u, decimal value)
        {
            setExact(divisor(u), value);
        }
        #endregion

        #region pretty formatting
        public Unit BestUnit
        {
            get
            {
                Unit u = Unit.B;
                while (u < Unit.GB)
                {
                    if (InUnitsExact(u) < 1000M)
                        break;
                    u++;
                }
                return u;
            }
        }
        public override string ToString()
        {
            Unit u = BestUnit;
            decimal d = InUnitsExact(u);
            return (Math.Round(d, 2)).ToString() + " " + UnitSuffixes[(ushort)u];
        }
        public static FileSize Parse(string s)
        {
            string[] parts = s.Split(' ');
            return new FileSize((Unit)Array.IndexOf(UnitSuffixes, parts[1]), decimal.Parse(parts[0]));

        }
        #endregion

        #region byte, KB, MB, GB accessors (exact and rounded)
        /// <summary>
        /// The size in bytes
        /// </summary>
        public ulong Bytes
        {
            get { return sizeInBytes; }
            set { sizeInBytes = value; }
        }

        /// <summary>
        /// The size in kilobytes (rounded)
        /// </summary>
        public ulong KB
        {
            get { return divideAndRound(1024); }
            set { checked { sizeInBytes = value * 1024; } }
        }

        /// <summary>
        /// The size in kilobytes (exact)
        /// </summary>
        public decimal KBExact
        {
            get { return ((decimal)sizeInBytes) / 1024M; }
            set { setExact(1024, value); }
        }

        /// <summary>
        /// The size in megabytes (rounded)
        /// </summary>
        public ulong MB
        {
            get { return divideAndRound(1024 * 1024); }
            set { checked { sizeInBytes = value * 1024 * 1024; } }
        }

        /// <summary>
        /// The size in megabytes (exact)
        /// </summary>
        public decimal MBExact
        {
            get { return ((decimal)sizeInBytes) / (1024M * 1024M); }
            set { setExact(1024 * 1024, value); }
        }

        /// <summary>
        /// Returns the size in gigabytes (rounded)
        /// </summary>
        public ulong GB
        {
            get { return divideAndRound(1024 * 1024 * 1024); }
            set { checked { sizeInBytes = value * 1024 * 1024 * 1024; } }
        }

        /// <summary>
        /// The size in gigabytes (exact)
        /// </summary>
        public decimal GBExact
        {
            get { return ((decimal)sizeInBytes) / (1024M * 1024M * 1024M); }
            set { setExact(1024 * 1024 * 1024, value); }
        }
        #endregion

        #region operators: + - * / < > <= >= == !=
        public static FileSize operator +(FileSize a, FileSize b)
        {
            FileSize result;
            checked
            {
                result.sizeInBytes = a.sizeInBytes + b.sizeInBytes;
            }
            return result;
        }
        public static FileSize operator -(FileSize a, FileSize b)
        {
            FileSize result;
            checked
            {
                result.sizeInBytes = a.sizeInBytes - b.sizeInBytes;
            }
            return result;
        }
        public static decimal operator /(FileSize a, FileSize b)
        {
            return ((decimal)a.sizeInBytes) / ((decimal)b.sizeInBytes);
        }

        public static FileSize operator /(FileSize a, decimal b)
        {
            return (a * (1.0M / b));
        }

        public static FileSize operator *(decimal a, FileSize b)
        {
            FileSize result;
            checked
            {
                result.sizeInBytes = (ulong)Math.Round(a * (decimal)b.sizeInBytes);
            }
            return result;
        }
        public static FileSize operator *(FileSize a, decimal b)
        {
            return b * a;
        }

        public static bool operator <(FileSize a, FileSize b)
        {
            return (a.sizeInBytes < b.sizeInBytes);
        }

        public static bool operator >(FileSize a, FileSize b)
        {
            return b < a;
        }

        public static bool operator ==(FileSize a, FileSize b)
        {
            return (a.sizeInBytes == b.sizeInBytes);
        }

        public static bool operator !=(FileSize a, FileSize b)
        {
            return !(a == b);
        }

        public static bool operator <=(FileSize a, FileSize b)
        {
            return (a.sizeInBytes <= b.sizeInBytes);
        }

        public static bool operator >=(FileSize a, FileSize b)
        {
            return (b <= a);
        }
        #endregion

        public override int GetHashCode()
        {
            return sizeInBytes.GetHashCode();
        }

        public override bool Equals(object obj)
        {
            return (obj is FileSize && ((FileSize)obj).sizeInBytes == sizeInBytes);
        }

        public FileSize(ulong numBytes)
        {
            sizeInBytes = numBytes;
        }

        public FileSize(long numBytes)
        {
            if (numBytes > 0)
                sizeInBytes = (ulong)numBytes;
            else
                sizeInBytes = 0;
        }

        public FileSize(Unit u, decimal value)
        {
            sizeInBytes = 0; // Dummy
            SetWithUnits(u, value);
        }
    }
    public class AudioBitrateCalculationStream
    {
        public string Filename;
        public FileSize? Size;
        public AudioType AType;
        public AudioBitrateCalculationStream(string filename)
        {
            this.Filename = filename;
            Size = FileSize.Of2(Filename);
            AType = AudioType.NONE;
            if (Path.GetExtension(Filename).ToLower() == ".ac3") AType = AudioType.AC3;
            else if (Path.GetExtension(Filename).ToLower() == ".aac") AType = AudioType.MP4AAC;
            else if (Path.GetExtension(Filename).ToLower() == ".mp3") AType = AudioType.MP3;
            else if (Path.GetExtension(Filename).ToLower() == ".mp2") AType = AudioType.MP3;
            else if (Path.GetExtension(Filename).ToLower() == ".dts") AType = AudioType.DTS;
            else AType = AudioType.AC3;
        }
    }
    public class CalcData
    {
        private static readonly decimal mp4OverheadWithBframes = 10.4M;
        private static readonly decimal mp4OverheadWithoutBframes = 4.3M;
        private static readonly decimal aviVideoOverhead = 24M;
        private static readonly decimal cbrMP3Overhead = 23.75M;
        private static readonly decimal vbrMP3Overhead = 40M;
        private static readonly decimal ac3Overhead = 23.75M;
        private static readonly int AACBlockSize = 1024;
        private static readonly int AC3BlockSize = 1536;
        private static readonly int MP3BlockSize = 1152;
        private static readonly int VorbisBlockSize = 1024;
        private static readonly int mkvAudioTrackHeaderSize = 140;
        private static readonly int mkvVorbisTrackHeaderSize = 4096;
        private static readonly uint mkvIframeOverhead = 26;
        private static readonly uint mkvPframeOverhead = 13;
        private static readonly uint mkvBframeOverhead = 16;
        private static readonly string qualityCodecModifierValues = "MPEG2=1.022,ASP=1.018,AVC=1.0,VC1=1.004";

        public CalcData(long frames, decimal fps) :
            this(frames, fps, ContainerType.MKV, null, true, null) { }

        public CalcData(long frames, decimal fps, ContainerType container, VideoCodec codec, bool bframes, AudioBitrateCalculationStream[] audios)
        {
            if (fps <= 0) throw new ArgumentException("Frames per second must be greater than zero.", "fps");
            if (frames <= 0) throw new ArgumentException("Frames must be greater than zero.", "frames");

            Frames = frames;
            FramesPerSecond = fps;
            ContainerType = container;
            VideoCodec = codec;
            HasBFrames = bframes;
            AudioStreams = audios;
            VideoOverheadRatio = 1;
            AudioOverheadRatio = 1;
            ExtraOverheadRatio = 1;
            QualityCoeffient = 0.75F;
        }

        public ContainerType ContainerType { get; set; }
        public VideoCodec VideoCodec { get; set; }
        public AudioBitrateCalculationStream[] AudioStreams { get; set; }
        public bool HasBFrames { get; set; }
        public decimal FramesPerSecond { get; protected set; }
        public long Frames { get; protected set; }
        public Size FrameSize { get; set; }
        public float BitsPerPixel { get; set; }

        public float QualityCodecModifier { get; set; }
        public float QualityCoeffient { get; set; }
        public float QualityEstimate { get; set; }

        public FileSize VideoOverhead { get; set; }
        public FileSize AudioOverhead { get; set; }
        public FileSize ExtraOverhead { get; set; }

        public float VideoOverheadRatio { get; set; }
        public float AudioOverheadRatio { get; set; }
        public float ExtraOverheadRatio { get; set; }

        public FileSize ExtraSize { get; set; }
        public FileSize VideoSize { get; set; }
        public FileSize AudioSize { get; set; }
        public FileSize TotalSize { get; set; }

        public decimal VideoBitrate
        {
            get { return (decimal)VideoSize.Bytes / 1000M / TotalSeconds * 8M; }
            set { VideoSize = new FileSize((ulong)(TotalSeconds * value * 1000M / 8M)); }
        }

        public decimal TotalSeconds { get { return (decimal)Frames / (decimal)FramesPerSecond; } }
        public FileSize VideoMuxSize { get { return new FileSize((ulong)((float)VideoSize.Bytes * VideoOverheadRatio) + VideoOverhead.Bytes); } }
        public FileSize AudioMuxSize { get { return new FileSize((ulong)((float)AudioSize.Bytes * AudioOverheadRatio) + AudioOverhead.Bytes); } }
        public FileSize ExtraMuxSize { get { return new FileSize((ulong)((float)ExtraSize.Bytes * ExtraOverheadRatio) + ExtraOverhead.Bytes); } }

        public void CalcByTotalSize()
        {
            CalcVideoOverhead();
            CalcAudioOverheadAndSize();
            CalcExtraOverhead();
            CalcQualityCodecModifier();
            VideoSize = new FileSize((ulong)((float)TotalSize.Bytes / VideoOverheadRatio)) - AudioMuxSize - ExtraMuxSize - VideoOverhead;
            CalcBitsPerPixel();
            CalcQualityEstimate();
        }

        public void CalcByVideoSize()
        {
            CalcVideoOverhead();
            CalcAudioOverheadAndSize();
            CalcExtraOverhead();
            CalcQualityCodecModifier();
            TotalSize = VideoMuxSize + AudioMuxSize + ExtraMuxSize;
            CalcBitsPerPixel();
            CalcQualityEstimate();
        }

        public void CalcByBitsPerPixel()
        {
            CalcVideoOverhead();
            CalcAudioOverheadAndSize();
            CalcExtraOverhead();
            CalcQualityCodecModifier();
            VideoSize = new FileSize((ulong)((
                BitsPerPixel / 8F * (float)(FrameSize.Width * FrameSize.Height) * Frames / VideoOverheadRatio)) - VideoOverhead.Bytes);
            TotalSize = VideoMuxSize + AudioMuxSize + ExtraMuxSize;
            CalcQualityEstimate();
        }

        public void CalcByQualityEstimate()
        {
            CalcVideoOverhead();
            CalcAudioOverheadAndSize();
            CalcExtraOverhead();
            CalcQualityCodecModifier();
            VideoBitrate = (decimal)QualityEstimate * ((decimal)Math.Pow((float)(FrameSize.Width * FrameSize.Height), QualityCoeffient * QualityCodecModifier) * FramesPerSecond) / 1000M;
            TotalSize = VideoMuxSize + AudioMuxSize + ExtraMuxSize;
            CalcBitsPerPixel();
        }

        private void CalcBitsPerPixel()
        {
            BitsPerPixel = (float)VideoSize.Bytes * 8F / (float)Frames / (float)(FrameSize.Width * FrameSize.Height);
        }

        private void CalcQualityEstimate()
        {
            QualityEstimate = (float)VideoBitrate / ((float)Math.Pow((float)(FrameSize.Width * FrameSize.Height), QualityCoeffient * QualityCodecModifier) * (float)FramesPerSecond) * 1000F;
        }

        private void CalcQualityCodecModifier()
        {
            var qualityCodecModifiers = new Dictionary<string, float>();

            // read the values into the dictionary
            foreach (string mod in qualityCodecModifierValues.Split(','))
            {
                qualityCodecModifiers.Add(mod.Split('=')[0], float.Parse(mod.Split('=')[1], System.Globalization.CultureInfo.InvariantCulture));
            }

            // use values when found in dictionary, otherwise default to no-modification
            if (qualityCodecModifiers.ContainsKey(VideoCodec.ID))
                QualityCodecModifier = qualityCodecModifiers[VideoCodec.ID];
            else
                QualityCodecModifier = 1F;
        }

        private void CalcVideoOverhead()
        {
            VideoOverheadRatio = 1;

            if (ContainerType == ContainerType.MP4)
            {
                VideoOverhead = new FileSize(Unit.B,
                    (HasBFrames ? mp4OverheadWithBframes : mp4OverheadWithoutBframes) * Frames);
            }
            else if (ContainerType == ContainerType.MKV)
            {
                long nbIframes = Frames / 10;
                long nbBframes = HasBFrames ? (Frames - nbIframes) / 2 : 0;
                long nbPframes = Frames - nbIframes - nbBframes;
                VideoOverhead = new FileSize(Unit.B,
                    (4300M + 1400M + nbIframes * mkvIframeOverhead + nbPframes * mkvPframeOverhead +
                    nbBframes * mkvBframeOverhead +
                    TotalSeconds * 12 / 2 // this line for 12 bytes per cluster overhoad
                    ));
            }
            else if (ContainerType == ContainerType.AVI)
            {
                VideoOverhead = new FileSize(Unit.B, Frames * aviVideoOverhead);
            }
            else if (ContainerType == ContainerType.M2TS)
            {
                VideoOverhead = new FileSize(Unit.B, GetM2tsVideoOverhead(AudioStreams.Length));
            }
        }

        /// <summary>
        /// gets the video overhead in the m2ts container
        /// </summary>
        /// <param name="iAudioStreamCount">audio stream count</param>
        /// <returns>overhead this video track will incurr</returns>
        private long GetM2tsVideoOverhead(int iAudioStreamCount)
        {
            long videoOverhead = 0;
            long videoSize = (long)VideoBitrate * (long)TotalSeconds; //videoOverhead+generalOverhead+rawVideo
            if (iAudioStreamCount == 0)
            {
                videoOverhead = Frames * 263 + 9412;
            }
            videoSize -= videoOverhead;
            int packageSize = 192;
            int packageFill = 184; //of 192
            long packageCount = videoSize / packageSize + 1;
            long packageOverhead = packageCount * (packageSize - packageFill);
            videoOverhead += packageOverhead;
            return videoOverhead;
        }

        private void CalcAudioOverheadAndSize()
        {
            AudioOverheadRatio = 1;
            AudioSize = FileSize.Empty;
            AudioOverhead = FileSize.Empty;

            foreach (var audio in AudioStreams)
            {
                AudioSize += (audio.Size ?? FileSize.Empty);

                if (ContainerType == ContainerType.MKV)
                    AudioOverhead += new FileSize(Unit.B, GetMkvAudioOverhead(audio.AType, 48000, (double)TotalSeconds));
                else if (ContainerType == ContainerType.M2TS)
                    AudioOverhead += new FileSize(Unit.B, GetM2tsAudioOverhead(audio.AType, Frames, (double)TotalSeconds, audio.Size.Value.Bytes, true));
                else if (ContainerType == ContainerType.AVI)
                    AudioOverhead += new FileSize(Unit.B, GetAviAudioOverhead(audio.AType) * Frames);
            }
        }

        private void CalcExtraOverhead()
        {
            if (ContainerType == ContainerType.M2TS)
                ExtraOverheadRatio = 2;
            else
                ExtraOverheadRatio = 1;
        }

        /// <summary>
        /// gets the overhead a given audio type will incurr in the matroska container
        /// given its length and sampling rate
        /// </summary>
        /// <param name="AudioType">type of the audio track</param>
        /// <param name="samplingRate">sampling rate of the audio track</param>
        /// <param name="length">length of the audio track</param>
        /// <returns>overhead this audio track will incurr</returns>
        private static int GetMkvAudioOverhead(AudioType audioType, int samplingRate, double length)
        {
            if (audioType == AudioType.NONE)
                return 0;
            Int64 nbSamples = Convert.ToInt64((double)samplingRate * length);
            int headerSize = mkvAudioTrackHeaderSize;
            int samplesPerBlock = 0;
            if (audioType == AudioType.MP4AAC || audioType == AudioType.M4A || audioType == AudioType.RAWAAC)
                samplesPerBlock = AACBlockSize;
            else if (audioType == AudioType.VBRMP3 || audioType == AudioType.CBRMP3 || audioType == AudioType.MP3 || audioType == AudioType.DTS)
                samplesPerBlock = MP3BlockSize;
            else if (audioType == AudioType.AC3)
                samplesPerBlock = AC3BlockSize;
            else if (audioType == AudioType.VORBIS)
            {
                samplesPerBlock = VorbisBlockSize;
                headerSize = mkvVorbisTrackHeaderSize;
            }
            else // unknown types..
            {
                samplesPerBlock = AC3BlockSize;
            }
            double blockOverhead = (double)nbSamples / (double)samplesPerBlock * 22.0 / 8.0;
            int overhead = (int)(headerSize + 5 * length + blockOverhead);
            return overhead;
        }

        /// <summary>
        /// gets the overhead a given audio type will incurr in the m2ts container
        /// </summary>
        /// <param name="AudioType">type of the audio track</param>
        /// <param name="lFrames">frames of the video track</param>
        /// <param name="length">length of the audio track</param>
        /// <param name="lFileSize">size of the audio track</param>
        /// <param name="bVideoHandling">true if video will also be muxed</param>
        /// <returns>overhead this audio track will incurr</returns>
        private static int GetM2tsAudioOverhead(AudioType audioType, long lFrames, double length, ulong lFileSize, bool bVideohandling)
        {
            int blocksize = 80;
            if (audioType == AudioType.AC3 || audioType == AudioType.THD || audioType == AudioType.DTS
             || audioType == AudioType.DTSHD || audioType == AudioType.DTSMA || audioType == AudioType.PCM)
            {
                blocksize = 907;
            }
            int audioOverhead = (int)((2 * blocksize) / 180.0 * 8.0 + 17.0 + 0.5) * ((int)(lFileSize * 1.0 / (blocksize * 1.0) + 0.5));
            if (bVideohandling)
            {
                double audiobitrate = lFileSize / length;
                audioOverhead += (int)lFrames * (262 + ((int)audiobitrate / 184 + 1)) + 9412;
            }
            return audioOverhead;
        }

        /// <summary>
        /// gets the avi container overhead for the given audio type and bitrate mode
        /// bitrate mode only needs to be taken into account for MP3 but it's there for all cases nontheless
        /// </summary>
        /// <param name="AudioType">the type of audio</param>
        /// <param name="bitrateMode">the bitrate mode of the given audio type</param>
        /// <returns>the overhead in bytes per frame</returns>
        private static decimal GetAviAudioOverhead(AudioType audioType)
        {
            if (audioType == AudioType.AC3)
                return ac3Overhead;
            else if (audioType == AudioType.MP3)
                return vbrMP3Overhead;
            else if (audioType == AudioType.VBRMP3)
                return vbrMP3Overhead;
            else if (audioType == AudioType.CBRMP3)
                return cbrMP3Overhead;
            else if (audioType == AudioType.RAWAAC)
                return vbrMP3Overhead;
            else if (audioType == AudioType.DTS)
                return ac3Overhead;
            else
                return 0;
        }
    }
}
